import torch
from mmengine.evaluator import BaseMetric


class IoU(BaseMetric):
    def process(self, data_batch, data_samples):
        preds, labels = data_samples[0], data_samples[1]["labels"]
        preds = torch.argmax(preds, dim=1)
        intersect = (labels == preds).sum()
        union = (torch.logical_or(preds, labels)).sum()
        iou = (intersect / union).cpu()
        self.results.append(dict(batch_size=len(labels), iou=iou * len(labels)))

    def compute_metrics(self, results):
        total_iou = sum(result["iou"] for result in self.results)
        num_samples = sum(result["batch_size"] for result in self.results)
        return dict(iou=total_iou / num_samples)
